from typing import Any

from vnpy.trader.object import TradeData
from vnpy_ctastrategy.backtesting import BacktestingEngine

from vnpy_ctastrategy import (
    CtaTemplate,
    BarGenerator,
    ArrayManager,
    TickData,
    BarData,
)


class CuatroStrategy(CtaTemplate):
    """
    此策略经过测试不大适用于螺纹钢，比较适用于股指期货
    盘整期 RSI 是在高出开空仓，趋势下 RSI 是在高开多仓。
    """
    author = "Cjie Wong"

    boll_window = 20
    boll_dev = 2.0
    rsi_window = 14
    rsi_signal = 30
    fast_window = 5
    slow_window = 20
    trailing_long = 1.0
    trailing_short = 1.0
    fixed_size = 1

    boll_up = 0
    boll_down = 0
    rsi_value = 0
    rsi_long = 0
    rsi_short = 0
    fast_ma = 0
    slow_ma = 0
    ma_trend = 0
    intra_trade_high = 0
    intra_trade_low = 0
    long_stop = 0
    short_stop = 0

    parameters = [
        "boll_window",
        "boll_dev",
        "rsi_window",
        "rsi_signal",
        "fast_window",
        "slow_window",
        "trailing_long",
        "trailing_short",
        "fixed_size"
    ]

    variables = [
        "boll_up",
        "boll_down",
        "rsi_value",
        "rsi_long",
        "rsi_short",
        "fast_ma",
        "slow_ma",
        "ma_trend",
        "intra_trade_high",
        "intra_trade_low",
        "long_stop",
        "short_stop"
    ]

    def __init__(self, cta_engine: Any, strategy_name: str, vt_symbol: str, setting: dict) -> None:
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)

        self.is_back_test = isinstance(cta_engine, BacktestingEngine)
        self.rsi_long = 50 + self.rsi_signal
        self.rsi_short = 50 - self.rsi_signal

        self.bg5 = BarGenerator(self.on_bar, 5, self.on_5min_bar)
        self.bg15 = BarGenerator(self.on_bar, 15, self.on_15min_bar)

        self.am5 = ArrayManager()
        self.am15 = ArrayManager()

    def on_init(self) -> None:
        self.write_log("策略初始化")
        self.load_bar(10)

    def on_start(self) -> None:
        self.write_log("策略启动")

    def on_stop(self) -> None:
        self.write_log("策略停止")

    def on_tick(self, tick: TickData) -> None:
        self.bg5.update_tick(tick)

    def on_bar(self, bar: BarData) -> None:
        if not self.is_back_test:
            print(bar)
        self.bg5.update_bar(bar)
        self.bg15.update_bar(bar)

    def on_5min_bar(self, bar: BarData):
        """"""
        self.cancel_all()

        self.am5.update_bar(bar)
        if not self.am5.inited or not self.am15.inited:
            return
        self.boll_up, self.boll_down = self.am5.boll(self.boll_window, self.boll_dev)
        self.rsi_value = self.am5.rsi(self.rsi_window)
        boll_width = self.boll_up - self.boll_down
        # boll_width = self.am5.std(self.boll_window) * self.boll_dev * 2

        # no position
        if self.pos == 0:
            self.intra_trade_high = bar.high_price
            self.intra_trade_low = bar.low_price
            self.long_stop = 0
            self.short_stop = 0

            if self.ma_trend > 0 and self.rsi_value >= self.rsi_long:
                self.buy(self.boll_up, self.fixed_size, stop=True)

            if self.ma_trend < 0 and self.rsi_value <= self.rsi_short:
                self.short(self.boll_down, self.fixed_size, stop=True)

        # Long position
        elif self.pos > 0:
            self.intra_trade_high = max(self.intra_trade_high, bar.high_price)
            self.long_stop = (self.intra_trade_high - self.trailing_long * boll_width)
            self.sell(self.long_stop, abs(self.pos), stop=True)

        # Short position
        else:
            self.intra_trade_low = min(self.intra_trade_low, bar.low_price)
            self.short_stop = (self.intra_trade_low + self.trailing_short * boll_width)
            self.cover(self.short_stop, abs(self.pos), stop=True)

    def on_15min_bar(self, bar: BarData):
        """"""
        self.am15.update_bar(bar)
        if not self.am15.inited:
            return
        self.fast_ma = self.am15.sma(self.fast_window)
        self.slow_ma = self.am15.sma(self.slow_window)

        if self.fast_ma > self.slow_ma:
            self.ma_trend = 1
        elif self.fast_ma < self.slow_ma:
            self.ma_trend = -1
        else:
            self.ma_trend = 0

        self.put_event()

    def on_trade(self, trade: TradeData) -> None:
        self.put_event()
